# -*- coding: utf-8 -*-
from __future__ import print_function
from ddl_platform.database.base import BaseTable
from ddl_platform.database.base import main_database
import pickle


class ResourceTable(BaseTable):
    def __init__(self, db=main_database):
        self.db = db
        self.TABLE = 'resourcetb'

    def query(self):
        resources = self.db.hvals(self.TABLE)
        decoded_r = [pickle.loads(r) for r in resources]
        return decoded_r

    def insert(self, resource):
        encoded_r = pickle.dumps(resource)
        self.db.hset(self.TABLE, resource.resource_id, encoded_r)
        return resource.resource_id 

    def update_available_mem(self, resource_id, mem):
        r = self.db.hval(self.TABLE, resource_id)
        r = pickle.loads(r)
        r.available_mem = mem
        encoded_r = pickle.dumps(r)
        self.db.hset(self.TABLE, resource.resource_id, encoded_r)

    def delete(self, resource_id):
        self.db.hdel(self.TABLE, resource_id)

    def clear_table(self):
        self.db.delete(self.TABLE)


tb = ResourceTable()
resource_table = tb

def clear_table():
    tb.clear_table()

def create_new_table():
    import ddl_platform.scheduler.resource_info as resource
    available_nodes = ['gpu1', 'gpu2', 'gpu4', 'gpu7', 'gpu8', 'gpu9', 'gpu10', 'gpu11', 'gpu12', 'gpu13', 'gpu14', 'gpu15', 'gpu16']
    gpu_ids = [0, 1, 2, 3]
    total_mem_per_gpu = 11*1024*1024*1024 # bytes
    for nid, node in enumerate(available_nodes):
        for gid in gpu_ids:
            r = resource.ResourceInfo(gid, total_mem_per_gpu, nid, node)
            tb.insert(r)
    print('total GPUs: ', len(available_nodes) * len(gpu_ids))

def test():
    entries = tb.query()
    print('# of entries: ', len(entries))


if __name__ == '__main__':
    #clear_table()
    test()
    #create_new_table()
